/*
  dispatcher:
  基于 JSON RPC 提供两个 API：
  1. QueryJudgeProcess：返回每个 judger 实例的运行状态
  2. DispatchSolution：从 backend 接收用户提交的代码，并下发给 judger 运行
*/

#include <stdio.h>
#include <string.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/types.h>
#include <sys/wait.h>

#include <map>
#include <memory>
#include <queue>
#include <sstream>
#include <stdexcept>
#include <string>

#include "common.h"
#include "config.h"
#include "helper.h"
#include "language.h"
#include "logger.h"
#include "mjson.h"  // use v1.2.7
#include "mjson.h"
#include "mongoose.h"  // use v1.1

/*
    dispatcher 的各种业务逻辑
*/
int mq_id;  // 消息队列 id

std::queue<DispatchSolutionReq> task_queue;
std::map<int, JudgeStatusMsg> judger_status_map;
std::map<int, JudgeStatusMsg> solution_status_map;

int gRunnerIdCounter = 1;  // 全局自增的 runner ID 计数器，用到天荒地老才会溢出

int recvAllMessage() {

  MsgBuffer msg;
  // 把消息都取出来
  int nbyte = 0;
  while ((nbyte = msgrcv(mq_id, (void *)&msg, sizeof(msg.mtext), MSG_TYPE, MSG_NOERROR | IPC_NOWAIT)) >= 0) {
    LOGD("recv msg %d bytes: %s\n", nbyte, msg.mtext);
    char *jstr = msg.mtext;

    JudgeStatusMsg st = json::parse(jstr, jstr + nbyte);

    if (judger_status_map.find(st.RunnerId) == judger_status_map.end()) {
      LOG_ERROR("get msg from a unexist judger %d", st.RunnerId);
    } else if (judger_status_map[st.RunnerId].IsJudgeOver == true) {
      LOG_ERROR("get msg from a stoppted judger %d", st.RunnerId);
    } else {
      LOGD("set judger rid %d pid %d sid %d tid %d", st.RunnerId, st.ProblemId, st.SolutionId, st.PassedTestCount);
      judger_status_map[st.RunnerId] = st;

      if (st.RunnerId <= 0) {
        LOG_FATAL("judger rid <= 0 !");
      }
    }
  }
  if (nbyte < 0) {
    // 失败处理，参考 man msgrcv 的 Program source
    if (errno != ENOMSG) {
      LOG_FATAL("msgrecv %s", strerror(errno));
      return -1;
    } else {
      // LOG_ERROR("ENOMSG");
    }
  }
  return 0;
}

void reclaimJudgerInstance() {
  for (auto iter = judger_status_map.begin(); iter != judger_status_map.end();) {
    // LOGD("iter judger rid %d", iter.first);
    auto &judge = (*iter).second;

    if (judge.RunnerId == 0) {
      LOG_FATAL("get a rid == 0 judger!!!, mp size %ld", judger_status_map.size());
    }

    // 更新 solution_status_map

    // TODO: 结构体转换可以封装一下，比如直接结构体嵌套就没这多事情了
    JudgeStatusMsg sol = judge;
    if (solution_status_map.find(sol.SolutionId) == solution_status_map.end()) {
      LOG_ERROR("get msg from a unexist solution %d", sol.SolutionId);
      LOGD("judge rid %d", judge.RunnerId);
    } else if (solution_status_map[sol.SolutionId].IsJudgeOver) {
      LOG_ERROR("get msg from a judge over solution %d", sol.SolutionId);
    } else {
      solution_status_map[sol.SolutionId] = sol;
    }

    // TODO: kill 10s 未更新的 judger

    // 已经结束的 judger，回收之
    // 已经结束的 solution 没法回收，因为不确定 backend 是否已经收到该评测结果
    if (judge.IsJudgeOver) {
      // C++ map 遍历过程中删除的陷阱
      iter = judger_status_map.erase(iter);
    } else {
      iter++;
    }
  }
}

int saveCodeFile(const DispatchSolutionReq &task, std::string &out_code_path) {
  std::string code_path;
  switch (task.Language) {
    case LANG_CPP:
      code_path = strFormat("%s/code/%d.cpp", OJ_HOME, task.SolutionId);
      break;
    case LANG_JAVA:
      code_path = strFormat("%s/code/%d.java", OJ_HOME, task.SolutionId);
      break;
    case LANG_PYTHON:
      code_path = strFormat("%s/code/%d.py", OJ_HOME, task.SolutionId);
      break;
    default:
      LOG_FATAL("unsupport language: %d", task.Language);
      exit(-1);
      break;
  }

  FILE *fp_code = fopen(code_path.c_str(), "we");
  if (fp_code == nullptr) {
    LOG_FATAL("create code %s failed, %s", code_path.c_str(), strerror(errno));
    return -1;
  }
  fprintf(fp_code, "%s", task.SourceCode.c_str());
  fclose(fp_code);

  LOGD("save code to %s", code_path.c_str());

  out_code_path = std::move(code_path);
  return 0;
}

int dispatchNewTask() {
  int ret = 0;

  while (task_queue.size() > 0 && judger_status_map.size() < JUDGER_INSTANCE_LIMIT) {
    DispatchSolutionReq task = task_queue.front();
    task_queue.pop();

    std::string user_code_path;
    ret = saveCodeFile(task, user_code_path);
    if (ret) {
      return ret;
    }

    // runner id 每次要记得 ++
    int runner_id = gRunnerIdCounter++;
    JudgeStatusMsg judger;
    judger.IsJudgeOver = false;
    judger.JudgeResult = oj_err_code_map[OJ_WT0];
    judger.LastUpdateTime = time(nullptr);
    judger.ProblemId = task.ProblemId;
    judger.RunnerId = runner_id;
    judger.SolutionId = task.SolutionId;
    judger.PassedTestCount = 0;
    judger.TotalTestCount = 0;
    judger.Language = task.Language;
    // 理论上不可能出现的情况
    if (judger_status_map.find(runner_id) != judger_status_map.end()) {
      LOG_FATAL("runner id %d dup !!!", runner_id);
      exit(-1);
    }
    // 新的 solution 和 judger 都必须在这里插入到 map
    // 后面的代码会根据是否在 map 里判断运行状态
    judger_status_map[runner_id] = judger;

    if (judger.RunnerId <= 0) {
      LOG_FATAL("judger rid <= 0 !");
    }

    pid_t pid = fork();
    if (pid == 0) {
      // 子进程负责开一个 judger 进程实例
      std::string rid_str = std::to_string(runner_id);
      std::string pid_str = std::to_string(task.ProblemId);
      std::string sid_str = std::to_string(task.SolutionId);
      std::string lang_str = std::to_string(task.Language);
      std::string judger_path = strFormat("%s/judger", OJ_HOME);

      // execl 要求以 null 结尾，否则会报 bad address
      ret = execl(judger_path.c_str(), "judger", "-r", rid_str.c_str(), "-p", pid_str.c_str(), "-s", sid_str.c_str(),
                  "-c", user_code_path.c_str(), "-l", lang_str.c_str(), nullptr);
      if (ret) {
        LOG_FATAL("execl failed, %s", strerror(errno));
        return -1;
      }
    } else {
      // 父进程打印下消息
      LOGD("create new judger, runner id %d process id %d sol id %d problem id %d", runner_id, pid, task.SolutionId,
           task.ProblemId);
    }
  }
  return 0;
}

/*
  mongoose HTTP 库使用参考 https://mongoose.ws/tutorials/http-server/
  mjson 做 JSON RPC 参考 https://www.mongoose.ws/tutorials/json-rpc-over-websocket/
  注意我们跟教程不同，因为我们不需要维护长连接，所以用的是 HTTP 协议，而不是 WebSocket 协议
*/

// constexpr 字符串参考 https://www.cnblogs.com/apocelipes/p/14769971.html

static void http_handler(struct mg_connection *c, int ev, void *ev_data, void *fn_data) {
  if (ev == MG_EV_HTTP_MSG) {
    // 所有的 HTTP 请求都视为 JSON RPC 请求处理
    mg_http_message *hm = (struct mg_http_message *)ev_data;
    mg_str req = hm->body;
    char *resp = nullptr;
    jsonrpc_process(req.ptr, req.len, mjson_print_dynamic_buf, &resp, nullptr);
    LOG(LL_INFO, ("[%.*s] -> [%s]", (int)req.len, req.ptr, resp));
    // 这里规定下，RPC 调用无论成功还是失败，返回的 HTTP 错误码都是 200
    // 具体的错误码应该写在 JSON 的 error 字段中
    // 参考 http://wiki.geekdream.com/Specification/json-rpc_2.0.html
    mg_http_reply(c, 200, "Content-Type: application/json\r\n", resp);
    free(resp);
  }
}

static void sum(struct jsonrpc_request *r) {
  double a = 0.0, b = 0.0;
  mjson_get_number(r->params, r->params_len, "$[0]", &a);
  mjson_get_number(r->params, r->params_len, "$[1]", &b);
  jsonrpc_return_success(r, "%g", a + b);
}

/*
request example:
{
  SolutionId: 123,
  ProblemId: 456
  SourceCode: "#include<stdio.h>int main(){ ....."
}
resp
{
  Accept: bool // 有空闲 judger 为 true，无空闲为 false
}
*/
static void DispatchSolution(struct jsonrpc_request *json_req) {
  LOGD("DispatchSolution req: %s", json_req->params);

  bool succ = false;
  bool accept = true;

  do {
    if (JUDGER_INSTANCE_LIMIT <= judger_status_map.size()) {
      // judger 实例已经耗尽
      accept = false;
      succ = true;  // 调用成功，但是实例耗尽，明示自己不接受任务
      break;
    }

    // 解析 task 对象
    DispatchSolutionReq task = json::parse(json_req->params, json_req->params + json_req->params_len);

    // 为了防止代码里有奇怪字符影响 JSON 解码，所以传输前做了 base64 加密
    std::string decode_str = b64decode(task.SourceCode.c_str(), task.SourceCode.length());
    task.SourceCode = std::move(decode_str);
    LOGD("decode: %s", task.SourceCode.c_str());

    // 1. 在 map 中初始化
    JudgeStatusMsg st;
    st.RunnerId = -1;
    st.SolutionId = task.SolutionId;
    st.PassedTestCount = 0;
    st.TotalTestCount = 0;
    st.JudgeResult = oj_err_code_map[OJ_WT0];
    st.IsJudgeOver = false;
    st.LastUpdateTime = time(0);
    st.Language = task.Language;
    solution_status_map[task.SolutionId] = st;

    // 2. push 到队列里
    task_queue.push(std::move(task));  // 这里用了 std move，所以别搞错 1 2 顺序

    accept = true;
    succ = true;

  } while (0);

  if (!succ) {
    LOGD("DispatchSolution resp: invalid args");
    // 错误码规范 http://wiki.geekdream.com/Specification/json-rpc_2.0.html
    jsonrpc_return_error(json_req, -32602, "Invalid params", nullptr);
    return;
  } else {
    DispatchSolutionResp resp;
    resp.Accept = accept;

    json j_resp;
    to_json(j_resp, resp);
    LOGD("DispatchSolution resp: %s", j_resp.dump().c_str());
    jsonrpc_return_success(json_req, j_resp.dump().c_str());
  }
}

/*
req
{
  SolutionId[] // 想要查询的 SolutionId
}
resp
{
  FreeJudgerNum: int
  JudgeStatus [
      {
          SolutionId: int
          RunnerId: int // 哪个 runner 负责
          ProblemId: int
          PassedTestCount: int // 当前通过几个测试点了
          TotalTestCount: int // 测试点总数
          JudgeResult: string // 最终评测结果
          IsJudgeOver: bool // 评测完了吗？
          LastUpdateTime: int
      }
  ]
}
*/
static void QueryJudgeProcess(struct jsonrpc_request *json_req) {
  // LOGD("QueryJudgeProcess req: %s", req->params);

  QueryJudgeProcessReq req = json::parse(json_req->params, json_req->params + json_req->params_len);

  QueryJudgeProcessResp resp;
  resp.FreeJudgerNum = (JUDGER_INSTANCE_LIMIT - judger_status_map.size());
  for (size_t i = 0; i < req.SolutionId.size(); i++) {
    int sid = req.SolutionId[i];
    if (solution_status_map.find(sid) != solution_status_map.end()) {
      // 调试的时候使用 JSON 格式化检验 https://www.bejson.com/
      auto &sol = solution_status_map[sid];
      resp.JudgeStatus.push_back(sol);
    }
  }

  // LOGD("QueryJudgeProcess resp:\n %s", ss.str().c_str());
  json j_resp = resp;
  jsonrpc_return_success(json_req, j_resp.dump().c_str());
}

int main(int argc, char *argv[]) {
  auto *ptr = getcwd(OJ_HOME, sizeof(OJ_HOME));
  if (ptr == nullptr) {
    LOG_FATAL("getcwd failed, %s", strerror(errno));
    exit(-1);
  }

  // Init manager
  struct mg_mgr mgr;
  mg_mgr_init(&mgr);

  // Init and register json rpc
  jsonrpc_init(nullptr, nullptr);
  jsonrpc_export("sum", sum);
  jsonrpc_export("DispatchSolution", DispatchSolution);
  jsonrpc_export("QueryJudgeProcess", QueryJudgeProcess);

  // Setup listener
  auto *conn_ptr = mg_http_listen(&mgr, LISTEN_URL, http_handler, &mgr);
  if (conn_ptr == nullptr) {
    LOG_FATAL("listen at %s failed!\n", LISTEN_URL);
    exit(-1);
  }
  LOGD("http server listen at %s\n", LISTEN_URL);

  // 开启 linux mq
  mq_id = msgget(MSG_KEY, IPC_CREAT | 0666);
  if (mq_id == -1) {
    LOG_FATAL("msgget fail, %s", strerror(errno));
    exit(EXIT_FAILURE);
  }

  for (;;) {
    // TODO: 这里的时间可以再思考下
    mg_mgr_poll(&mgr, 300);  // Event loop
    // LOGD("main logic call");
    int ret = recvAllMessage();
    if (ret) {
      LOG_FATAL("recvAllMessage err, ret %d", ret);
      exit(-1);
    }
    reclaimJudgerInstance();
    ret = dispatchNewTask();
    if (ret) {
      LOG_FATAL("dispatchNewTask err, ret %d", ret);
      exit(-1);
    }
  }

  mg_mgr_free(&mgr);  // Cleanup
  return 0;
}
